package top.yudoge.dogeeditor.llm.rewrite

import cn.hutool.core.util.StrUtil
import cn.hutool.json.JSONUtil
import org.springframework.stereotype.Component

@Component
@Deprecated(message = "This chain node is deprecated and will be removed in the future.")
class TableNameExtractChainNode : TaskChainNode {

    override fun process(ctx: ChainContext) {
        sendReasoningCallback(ctx, "guessing table names")

        val tableNames = mutableSetOf<String>()
        val query = ctx.query

        tableNames.addAll(ctx.query.tableSelected)

        tableNames.addAll(
            extractTableNamesFromCode(ctx, query.code)
        )

        tableNames.addAll(
            extractTableNamesFromUserTask(ctx, query.userTask)
        )

        println("${tableNames.size} table names extracted: $tableNames")
        putChainStore(ctx, TABLE_NAMES, tableNames)
        emitNext(ctx)
    }

    private fun extractTableNamesFromUserTask(ctx: ChainContext, userTask: String): Collection<String> {
        if (StrUtil.isBlank(userTask)) {
            return setOf()
        }
        var result = ctx.chatModel.chat(
            "These is a user task about SQL-script rewrite. Please extract the keywords that you think it is a table name. Please extract as more as possible. \n" +
                    "User task:\n" +
                    "$userTask\n\n" +
                    "Please return in the JSON format without any other information\n" +
                    "Example output: [\"table name 1\", \"table name 2\"]"
        )
        return extractModelOutput(result)
    }


    fun extractTableNamesFromCode(ctx: ChainContext, code: String): Set<String> {

        if (StrUtil.isBlank(code)) {
            return setOf()
        }

        var result = ctx.chatModel.chat(
            "Extracting the table name from given code: \n" +
                    " ```sql\n" +
                    "$code\n" +
                    "```\n" +
                    "Please return in the JSON format without any other information\n" +
                    "Example output: [\"table name 1\", \"table name 2\"]"
        )
        return extractModelOutput(result)
    }

    private fun extractModelOutput(output: String): Set<String> {
        return try {
            JSONUtil.toList(output, String::class.java).toSet()
        } catch (e: Exception) {
            setOf()
        }
    }

    companion object {
        const val TABLE_NAMES = "tableNames"
    }

}